package com.shoulder.boot.ops.ws;

import com.jcraft.jsch.ChannelShell;
import com.jcraft.jsch.Session;
import com.shoulder.boot.common.ws.WsSessionManager;
import com.shoulder.boot.ops.entity.Machine;
import com.shoulder.boot.ops.service.MachineService;
import com.shoulder.core.exception.BizException;
import com.shoulder.ssh.properties.SSHProperties;
import com.shoulder.ssh.utils.ShellUtils;
import lombok.Data;
import lombok.EqualsAndHashCode;
import lombok.NoArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;


@EqualsAndHashCode(callSuper = true)
@Slf4j
@Data
@NoArgsConstructor
public class SSHSocketHandler extends TextWebSocketHandler {

    private MachineService machineService;
    private Session session;
    private ChannelShell channel;
    private SSHProperties properties;
    private ShellUtils shellUtils;

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        Object token = session.getAttributes().get("token");
        if (token != null) {
            // 连接成功，放入在线缓存
            WsSessionManager.add(token.toString(), session);
        } else {
            throw new BizException("用户登录已经失效!");
        }
        Machine machine = machineService.noneScopeById(token.toString());
        if (machine == null) {
            throw new BizException("无效的终端设备");
        }
        shellUtils = new ShellUtils();
        shellUtils.setProperties(properties);
        this.session = shellUtils.init(machine.getHost(), machine.getPort(), machine.getUsername(), machine.getPassword());
        channel = shellUtils.initShell(this.session);
        shellUtils.startForwarding(channel, session);
    }


    @Override
    public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception {
        // 获得客户端传来的消息
        String payload = ((TextMessage) message).getPayload();
        Object token = session.getAttributes().get("token");
        log.info("SSHSocketHandler server 接收到 {} 发送的 {}", token, payload);

        shellUtils.out(channel, payload);

    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        log.info("afterConnectionClosed.......");
        Object token = session.getAttributes().get("token");
        if (token != null) {
            // 退出，移除缓存
            WsSessionManager.remove(token.toString());
        }
        if (channel != null) {
            channel.disconnect();
        }

        if (this.session != null) {
            this.session.disconnect();
        }

    }
}
